import os
from collections import defaultdict


def convert(distribution_path):
    with open(distribution_path, 'r', encoding='utf-8') as f:
        original_distribution = eval(f.read())
    tensor_info = {}
    other_info = defaultdict(lambda: defaultdict(int))

    t_name_to_shapes = {}
    t_infos = []
    o_infos = []
    for key in original_distribution:
        for t_info in key:
            t_name = t_info[0]
            t_shape = t_info[1]
            if t_name not in t_name_to_shapes:
                t_name_to_shapes[t_name] = set()
            t_name_to_shapes[t_name].add(t_shape)

        t_infos.append(original_distribution[key]['tensor_info'])
        o_infos.append(original_distribution[key]['other_info'])

    t_name_to_values = {}
    for t_info in t_infos:
        for t_name, val in t_info.items():
            if t_name not in t_name_to_values:
                t_name_to_values[t_name] = {}

            for t_dtype, t_distribution in val.items():
                cur_max = t_distribution['max']
                cur_min = t_distribution['min']

                if t_dtype not in t_name_to_values[t_name]:
                    t_name_to_values[t_name][t_dtype] = {
                        'max': cur_max,
                        'min': cur_min,
                    }
                else:
                    if cur_max > t_name_to_values[t_name][t_dtype]['max']:
                        t_name_to_values[t_name][t_dtype]['max'] = cur_max
                    if cur_min < t_name_to_values[t_name][t_dtype]['min']:
                        t_name_to_values[t_name][t_dtype]['min'] = cur_min

    for o_info in o_infos:
        for o_name, v_info in o_info.items():
            for val, count in v_info.items():
                other_info[o_name][val] += count

    for o_name, val_counts in other_info.items():
        total = sum(val_counts.values())
        for val, count in val_counts.items():
            other_info[o_name][val] = count / total

    for t_name in t_name_to_values:
        tensor_info[t_name] = {'value_distribution': t_name_to_values[t_name]}
        tensor_info[t_name]['shape'] = t_name_to_shapes[t_name]

    other_info = dict_recursive(dict(other_info))
    for t_name in tensor_info:
        other_info.pop(t_name, None)

    api_name = os.path.basename(distribution_path)
    store_distribution_info(api_name, tensor_info, other_info)


def dict_recursive(d):
    for k, v in d.items():
        if isinstance(v, defaultdict):
            d[k] = dict_recursive(dict(v))
        elif isinstance(v, dict):
            d[k] = dict_recursive(v)
    return d


def store_distribution_info(api_name, tensor_type_info, other_type_info, store_dir='./hwei'):
    with open(os.path.join(store_dir, api_name), 'w', encoding='utf-8') as f:
        api_distribution_info = {
            'tensor_info': tensor_type_info,
            'other_info': other_type_info
        }
        f.write(str(api_distribution_info))


if __name__ == '__main__':
    convert('./distribution3/torch.abs')
